package com.ztom.top100;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 前 K 个高频元素
 * <p>
 * https://leetcode-cn.com/problems/top-k-frequent-elements/
 *
 * @author ZhangTao
 */
public class Code83TopKFrequent {

    /**
     * 小根堆
     *
     * @param nums
     * @param k
     * @return
     */
    public int[] topKFrequent1(int[] nums, int k) {
        int[] res = new int[k];
        if (nums == null || nums.length == 0 || k <= 0) {
            return res;
        }
        Map<Integer, Integer> countMap = new HashMap<>();
        // 统计频次
        for (int num : nums) {
            countMap.put(num, countMap.getOrDefault(num, 0) + 1);
        }
        // 小根堆
        PriorityQueue<int[]> heap = new PriorityQueue<>(Comparator.comparingInt(a -> a[1]));
        countMap.forEach((num, count) -> {
            if (heap.size() < k) {
                heap.add(new int[]{num, count});
            } else {
                if (count > heap.peek()[1]) {
                    heap.poll();
                    heap.add(new int[]{num, count});
                }
            }
        });
        // 输出结果
        for (int i = 0; i < k; i++) {
            res[i] = heap.poll()[0];
        }
        return res;
    }

    /**
     * 改写快排
     *
     * @param nums
     * @param k
     * @return
     */
    public int[] topKFrequent(int[] nums, int k) {
        int[] res = new int[k];
        if (nums == null || nums.length == 0 || k <= 0) {
            return res;
        }

        // 统计频次
        Map<Integer, Integer> countMap = new HashMap<>();
        for (int num : nums) {
            countMap.put(num, countMap.getOrDefault(num, 0) + 1);
        }

        List<int[]> countList = new ArrayList<>();
        countMap.forEach((num, count) -> {
            countList.add(new int[]{num, count});
        });

        qsort(countList, res, k - 1);

        return res;
    }

    private void qsort(List<int[]> countList, int[] res, int k) {
        int l = 0;
        int r = countList.size() - 1;
        int[] mid = null;
        int pivot = 0;
        while (l <= r) {
            pivot = countList.get(l + (int) (Math.random() * (r - l + 1)))[1];
            mid = partition(countList, l, r, pivot);
            if (k > mid[1]) {
                l = mid[1] + 1;
            } else if (k < mid[0]) {
                r = mid[0] - 1;
            } else {
                // 收集答案
                for (int i = 0; i <= k; i++) {
                    res[i] = countList.get(i)[0];
                }
                return;
            }
        }
    }

    private int[] partition(List<int[]> countList, int l, int r, int pivot) {
        if (l >= r) {
            return new int[]{l, r};
        }
        int more = l - 1;
        int less = r + 1;
        int i = l;
        while (i < less) {
            if (countList.get(i)[1] == pivot) {
                i++;
            } else if (countList.get(i)[1] > pivot) {
                Collections.swap(countList, i++, ++more);
            } else {
                Collections.swap(countList, i, --less);
            }
        }
        return new int[]{more + 1, less - 1};
    }
}
